"""
qwen2指令微调模型-训练脚本
"""
import json
import pandas as pd
import torch
from datasets import Dataset
from modelscope import snapshot_download, AutoTokenizer
from swanlab.integration.huggingface import SwanLabCallback
from peft import LoraConfig, TaskType, get_peft_model
from transformers import AutoModelForCausalLM, TrainingArguments, Trainer, DataCollatorForSeq2Seq
import os
import swanlab

def dataset_jsonl_transfer(origin_path, new_path):
    """
    将原始数据集转换为大模型微调所需数据格式的新数据集
    """
    messages = []
    # 读取原始数据集JSONL文件
    with open(origin_path, "r") as file:
        for line in file:
            # 解析每一行的json数据
            data = json.loads(line)
            context = data["text"]
            catagory = data["category"]
            label = data["output"]
            message = {
                "instruction": "你是一个文本分类领域的专家，你会接收到一段文本和几个潜在的分类选项，请输出文本内容的正确类型",
                "input": f"文本:{context},类型选型:{catagory}",
                "output": label,
            }
            messages.append(message)

    # 保存重构后的JSONL文件
    with open(new_path, "w", encoding="utf-8") as file:
        for message in messages:
            file.write(json.dumps(message, ensure_ascii=False) + "\n")

def process_func(example):
    """
    将数据集进行预处理成指令微调需要用到的测试集。
    这段代码的主要作用是将【系统指令、用户输入】和【模型响应】组合成一个适合模型训练的【输入ID序列】，并生成相应的【注意力掩码序列】和【标签序列】。
    通过这种方式，模型可以学习根据输入文本和选项进行分类
    """
    # 定义输入序列的最大长度为384
    MAX_LENGTH = 384
    # 初始化三个空列表，用于存储输入ID、注意力掩码和标签
    input_ids, attention_mask, labels = [], [], []
    # 使用tokenizer将【系统指令、用户输入】编码为token ID，并且不添加特殊标记（如CLS、SEP等）
    instruction = tokenizer(
        f"<|im_start|>system\n你是一个文本分类领域的专家，你会接收到一段文本和几个潜在的分类选项，请输出文本内容的正确类型<|im_end|>\n<|im_start|>user\n{example['input']}<|im_end|>\n<|im_start|>assistant\n",
        add_special_tokens=False,
    )
    # 使用tokenizer将【模型响应】编码为token ID，同样不添加特殊标记
    response = tokenizer(f"{example['output']}", add_special_tokens=False)
    # 【输入ID序列】：将指令的input_ids、响应的input_ids以及一个垫定token ID（tokenizer.pad_token_id）拼接在一起，形成最终的input_ids
    input_ids = instruction["input_ids"] + response["input_ids"] + [tokenizer.pad_token_id]
    # 【注意力掩码序列】：将指令的attention_mask、响应的attention_mask以及一个1拼接在一起，形成最终的attention_mask。注意，垫定token的attention_mask设为1
    attention_mask = (
        instruction["attention_mask"] + response["attention_mask"] + [1]
    )
    """
    【标签序列】，其中：
    - 指令部分的标签设为-100（通常用于忽略不参与损失计算的部分）。
    - 响应部分的标签设为对应的input_ids。
    - 最后一个垫定token的标签设为垫定token ID。
    """
    labels = [-100] * len(instruction["input_ids"]) + response["input_ids"] + [tokenizer.pad_token_id]
    # 截断上述3个序列
    if len(input_ids) > MAX_LENGTH:
        input_ids = input_ids[:MAX_LENGTH]
        attention_mask = attention_mask[:MAX_LENGTH]
        labels = labels[:MAX_LENGTH]
    return {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels}

def predict(messages, model, tokenizer):
    """
    根据输入的消息生成模型的回复
    """
    # 使用GPU进行计算
    device = "cuda"
    # 使用tokenizer.apply_chat_template将messages格式化为适合模型输入的字符串text
    text = tokenizer.apply_chat_template(
        messages,
        # 输入是否要分词，这里不需要分词，直接输出文本
        tokenize=False,
        # 在文本末尾添加生成提示
        add_generation_prompt=True
    )
    """
    使用tokenizer将格式化后的text编码为PyTorch张量，并传入模型进行处理。
    return_tensors="pt"表示返回PyTorch张量。
    将张量移动到指定设备（GPU）上。
    """
    model_inputs = tokenizer([text], return_tensors="pt").to(device)
    # 调用model.generate生成回复的token ID。
    generated_ids = model.generate(
        model_inputs.input_ids,
        # 生成最多512个新的tokens
        max_new_tokens=512
    )
    # 通过列表推导式，从generated_ids中排除输入部分的token ID，只保留生成的新内容。
    generated_ids = [
        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
    ]
    """
    使用tokenizer.batch_decode将生成的token ID解码为文本。
    skip_special_tokens=True表示跳过特殊符号，使生成的文本更干净。
    """
    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]

    print(response)

    return response





# 1.下载Qwen模型和加载模型权重
# 在modelscope上下载Qwen模型到本地目录下
model_dir = snapshot_download("qwen/Qwen2-1.5B-Instruct", cache_dir="./", revision="master")

# Transformers加载模型权重（tokenizer和model）
tokenizer = AutoTokenizer.from_pretrained("./qwen/Qwen2-1___5B-Instruct/", use_fast=False, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained("./qwen/Qwen2-1___5B-Instruct/", device_map="auto", torch_dtype=torch.bfloat16)
# 开启梯度检查点时，要执行该方法
model.enable_input_require_grads()

# 2.加载、处理训练集和测试集
train_dataset_path = "train.jsonl"
test_dataset_path = "test.jsonl"

train_dataset_new_path = "new_train.jsonl"
test_dataset_new_path = "new_test.jsonl"

# 新训练集和测试集文件不存在时，需要将原始训练集和测试集转换为新训练集和测试集文件
if not os.path.exists(train_dataset_new_path):
    dataset_jsonl_transfer(train_dataset_path, train_dataset_new_path)
if not os.path.exists(test_dataset_new_path):
    dataset_jsonl_transfer(test_dataset_path, test_dataset_new_path)

# 得到预处理后的训练集
train_df = pd.read_json(train_dataset_new_path, lines=True)
train_ds = Dataset.from_pandas(train_df)
train_dataset = train_ds.map(process_func, remove_columns=train_ds.column_names)

# 3.训练模型
# Lora配置
config = LoraConfig(
    # 指定任务类型为因果语言模型
    task_type=TaskType.CAUSAL_LM,
    # 指定在哪些层应用LoRA，包括注意力机制和前馈网络的投影层。
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    # 训练模式
    inference_mode=False,
    # Lora 秩
    r=8,
    # Lora缩放因子，具体作用参见 Lora 原理
    lora_alpha=32,
    # Lora dropout 比例
    lora_dropout=0.1
)
# 将LoRA配置应用到模型，使其在指定层使用LoRA技术进行训练
model = get_peft_model(model, config)
# 训练参数设置
args = TrainingArguments(
    # 训练好的模型保存在这里
    output_dir="./output/Qwen2",
    # 每GPU批量大小
    per_device_train_batch_size=4,
    # 梯度累积步数，实际批量大小为4x4=16
    gradient_accumulation_steps=4,
    # 每10步记录日志
    logging_steps=10,
    # 训练 2 epoch，(也就是epoch最终=2)
    num_train_epochs=2,
    # 每100步保存模型
    save_steps=100,
    # 学习率，从10e-5开始逐步往下降，每10步降0.2e-5
    learning_rate=1e-4,
    # 每个节点保存模型
    save_on_each_node=True,
    # 开启梯度检查点以节省内存
    gradient_checkpointing=True,
    # 不向任何平台报告训练指标
    report_to="none"
)
# 定义SwanLabCallback回调函数，用于记录项目信息和训练配置
swanlab_callback = SwanLabCallback(
    project="Qwen2-fintune",
    experiment_name="Qwen2-1.5B-Instruct",
    description="使用通义千问Qwen2-1.5B-Instruct模型在zh_cls_fudan-news数据集上微调。",
    config={
        "model": "qwen/Qwen2-1.5B-Instruct",
        "dataset": "swift/zh_cls_fudan-news"
    }
)
# 创建Trainer实例，传入模型、训练参数、训练数据集、数据整理器DataCollatorForSeq2Seq和回调函数
trainer = Trainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer, padding=True),
    callbacks=[swanlab_callback],
)
# 开始训练
trainer.train()

# 4.测试模型
# 用测试集的前10条，测试模型
test_df = pd.read_json(test_dataset_new_path, lines=True)[:10]
# 开始测试模型
test_text_list = []
for index, row in test_df.iterrows():
    instruction = row['instruction']
    input_value = row['input']

    messages = [
        {"role": "system", "content": f"{instruction}"},
        {"role": "user", "content": f"{input_value}"}
    ]

    response = predict(messages, model, tokenizer)
    messages.append({"role": "assistant", "content": f"{response}"})
    result_text = f"{messages[0]}\n\n{messages[1]}\n\n{messages[2]}"
    test_text_list.append(swanlab.Text(result_text, caption=response))

# 使用swanlab.log({"Prediction": test_text_list})记录测试结果，其中test_text_list包含所有测试对话的文本内容和模型回复。
swanlab.log({"Prediction": test_text_list})
# 调用swanlab.finish()结束日志记录。
swanlab.finish()
